misc: Fix persistent kernel compilation #1430
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @Edenzzzz, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
I've addressed a persistent kernel compilation issue by re-introducing the normalize_d function. This function was inadvertently removed in a prior pull request (#1389) and is essential for the attention mechanism's normalization process, ensuring correct kernel compilation and functionality.
Highlights
- Re-introduction of
normalize_dfunction: Thenormalize_ddevice function, which was previously removed, has been re-introduced toinclude/flashinfer/attention/persistent.cuh. This function is crucial for normalizing output fragments (o_frag) using reciprocal values ofd, specifically when softmax is enabled in the attention variant.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
Edenzzzz
changed the title
There was a problem hiding this comment.
Code Review
This pull request restores the normalize_d function in include/flashinfer/attention/persistent.cuh, which was removed in a previous change. The implementation looks correct and addresses the compilation issue. I have one minor suggestion to improve the code's efficiency and readability by using bitwise operations for an index calculation.
| o_frag[mma_q][mma_d][reg_id] = | ||
| o_frag[mma_q][mma_d][reg_id] * d_rcp[mma_q][(reg_id % 4) / 2]; |
There was a problem hiding this comment.
The index calculation (reg_id % 4) / 2 can be expressed more efficiently and idiomatically using bitwise operations. Using bitwise operations can be slightly faster and is common practice in performance-critical CUDA code.
o_frag[mma_q][mma_d][reg_id] =
o_frag[mma_q][mma_d][reg_id] * d_rcp[mma_q][(reg_id >> 1) & 1];
📌 Description
#1389 removed
normalize_dfromprefill.cuh, now adding it backcc @yzh119
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes